CODE_PATH = '/home/xiaomin/wxm/Code/KaggleCCS'
caffe_root = '/opt/caffe/'  # this file should be run from {caffe_root}/examples (otherwise change this line)
data_root = '/home/xiaomin/wxm/Data/KaggleCCS/'
import sys
import numpy as np

sys.path.insert(0, caffe_root + 'python')
import caffe
import matplotlib.pyplot as plt
import os
import time

caffe.set_device(1)
caffe.set_mode_gpu()

model_def = CODE_PATH + '/' + 'prototxt/vgg/VGG_2014_16_deploy.prototxt'
model_weights = '/home/xiaomin/wxm/Data/KaggleCCS/snapshots/vgg2/_iter_2000.caffemodel'
batch_size = 16
test_image_path = '/home/xiaomin/wxm/Data/KaggleCCS/test'
submit_file = '/home/xiaomin/wxm/Data/KaggleCCS/submissions/vgg2/i2000.csv'

net = caffe.Net(model_def,  # defines the structure of the model
                model_weights,  # contains the trained weights
                caffe.TEST)  # use test mode (e.g., don't perform dropout)

# load the mean ImageNet image (as distributed with Caffe) for subtraction
mu = np.load(caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy')
mu = mu.mean(1).mean(1)  # average over pixels to obtain the mean (BGR) pixel values
print 'mean-subtracted values:', zip('BGR', mu)


def run_per_batch(transformed_data):
    net.blobs['data'].data[...] = transformed_data

    ### perform classification
    output = net.forward()

    output_prob = output['prob']  # the output probability vector for the first image in the batch
    return output_prob


# create transformer for the input called 'data'
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})

transformer.set_transpose('data', (2, 0, 1))  # move image channels to outermost dimension
transformer.set_mean('data', mu)  # subtract the dataset-mean value in each channel
transformer.set_raw_scale('data', 255)  # rescale from [0, 1] to [0, 255]
transformer.set_channel_swap('data', (2, 1, 0))  # swap channels from RGB to BGR

net.blobs['data'].reshape(batch_size,  # batch size
                          3,  # 3-channel (BGR) images
                          224, 224)  # image size is 227x227

f_submit = open(submit_file, 'w')
line = 'image_name' + ',' + 'Type_1' + ',' + 'Type_2' + ',' + 'Type_3' + '\n'
f_submit.write(line)


def per_im():
    nb = 0
    for im_name in os.listdir(test_image_path):
        start = time.time()
        image = caffe.io.load_image(os.path.join(test_image_path, im_name))
        transformed_image = transformer.preprocess('data', image)
        # plt.imsave('0.jpg', transformed_image)
        # plt.imshow(transformed_image)
        transformed_data = np.expand_dims(transformed_image, 0)

        output_prob = run_per_batch(transformed_data)
        line = im_name + ',' + str(output_prob[0]) + ',' + str(output_prob[1]) + ',' + str(output_prob[2]) + '\n'
        f_submit.write(line)
        end = time.time()

        nb += 1
        print 'nb.' + str(nb) + ':' + im_name + '  sec/im:' + str(end - start)


def per_batch(batch_size):
    im_names = []
    for im_name in os.listdir(test_image_path):
        im_names.append(im_name)

    nb = 0
    for i in range(len(im_names) / batch_size):
        start = time.time()
        names = []
        transformed_data_set = None
        for j in range(batch_size):
            names.append(im_names[nb])
            image = caffe.io.load_image(os.path.join(test_image_path, im_names[nb]))
            transformed_image = transformer.preprocess('data', image)
            # plt.imsave('0.jpg', transformed_image)
            # plt.imshow(transformed_image)
            transformed_data = np.expand_dims(transformed_image, 0)

            if transformed_data_set is None:
                transformed_data_set = transformed_data
            else:
                transformed_data_set = np.concatenate((transformed_data_set, transformed_data), 0)
            nb += 1

        output_prob = run_per_batch(transformed_data_set)
        for k in range(len(output_prob)):
            line = names[k] + ',' + str(output_prob[k][0]) + ',' + str(output_prob[k][1]) + ',' + str(
                output_prob[k][2]) + '\n'
            f_submit.write(line)
        end = time.time()

        print 'nb.' + str(nb) + '  sec/batch:' + str(end - start)


if __name__ == '__main__':
    per_batch(batch_size)
